package com.example.frameboot.common;

import com.alibaba.excel.EasyExcel;
import com.alibaba.excel.ExcelWriter;
import com.alibaba.excel.write.metadata.WriteSheet;
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import lombok.extern.slf4j.Slf4j;

import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.io.OutputStream;
import java.io.UnsupportedEncodingException;
import java.net.URLEncoder;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicLong;
import java.util.function.Function;

/**
 * Excel工具类
 * @author www@yiynx.cn
 */
@Slf4j
public class ExcelUtil {
    public static final Integer EXCEL_SHEET_ROW_MAX_SIZE = 1000001; // excel sheet最大行数(算标题)
    private static final Integer PAGE_SEARCH_TIMEOUT_SECONDS = 60;
    private static final int DEF_QUEUE_SIZE = Math.min(Runtime.getRuntime().availableProcessors(), 4);

    private HttpServletResponse httpServletResponse;
    private boolean parallel;
    private int queueSize;

    private ExcelUtil() {}

    public static ExcelUtil download(HttpServletResponse response, String excelFileName) throws UnsupportedEncodingException {
        response.setContentType("application/vnd.openxmlformats-officedocument.spreadsheetml.sheet");
        response.setCharacterEncoding("utf-8");
        String fileName = URLEncoder.encode(excelFileName, "UTF-8").replaceAll("\\+", "%20");
        response.setHeader("Content-disposition", "attachment;filename*=utf-8''" + fileName + ".xlsx");
        ExcelUtil excelUtil = new ExcelUtil();
        excelUtil.httpServletResponse = response;
        return excelUtil;
    }

    public ExcelUtil parallel() {
        this.parallel = true;
        this.queueSize = DEF_QUEUE_SIZE;
        return this;
    }

    public ExcelUtil parallel(int queueSize) {
        this.parallel = true;
        this.queueSize = queueSize;
        return this;
    }

    public void pageExcelWriter(Class head, Function<Page, Page> pageFunction) throws IOException {
        if (parallel) {
            pageExcelWriterQuick(head, queueSize, pageFunction);
        } else {
            pageExcelWriter(httpServletResponse.getOutputStream(), head, pageFunction);
        }
    }

    public static void pageExcelWriter(OutputStream outputStream, Class head, Function<Page, Page> pageFunction) {
        long start = System.currentTimeMillis();
        log.debug("excel writer start");
        try (ExcelWriter excelWriter = EasyExcel.write(outputStream, head).build()) {
            Page page = null;
            WriteSheet writeSheet = EasyExcel.writerSheet(1).build();
            do {
                long pageSearchStartTime = System.currentTimeMillis();
                page = pageFunction.apply(page == null ? new Page(1, 1000) : new Page(page.getCurrent() + 1, 1000)); // 分页查询
                long pageExcelWriteStartTime = System.currentTimeMillis();
                int sheetNo = (int) (page.getCurrent() * page.getSize() / EXCEL_SHEET_ROW_MAX_SIZE + 1);
                if (writeSheet.getSheetNo() != sheetNo) {
                    writeSheet = EasyExcel.writerSheet(sheetNo).build();
                }
                excelWriter.write(page.getRecords(), writeSheet); // excel写入数据
                log.debug("total:{}, pageSize:{}, totalPage:{}, pageNo:{}, sheetNo:{}, pageSearchTime:{}ms, pageExcelWriterTime:{}ms", page.getTotal(), page.getSize(), page.getPages(), page.getCurrent(), sheetNo, pageExcelWriteStartTime - pageSearchStartTime, System.currentTimeMillis() - pageExcelWriteStartTime);
            } while (page.hasNext()); // 是否还有下一页
        }
        log.debug("excel writer done, totalTime:{}ms", System.currentTimeMillis() - start);
    }

    private void pageExcelWriterQuick(Class head, Function<Page, Page> pageFunction) throws IOException {
        pageExcelWriterQuick(httpServletResponse.getOutputStream(), head, DEF_QUEUE_SIZE, pageFunction);
    }

    private void pageExcelWriterQuick(Class head, int queueSize, Function<Page, Page> pageFunction) throws IOException {
        pageExcelWriterQuick(httpServletResponse.getOutputStream(), head, queueSize, pageFunction);
    }

    public static void pageExcelWriterQuick(OutputStream outputStream, Class head, int queueSize, Function<Page, Page> pageFunction) {
        long start = System.currentTimeMillis();
        AtomicLong readTotalTime = new AtomicLong();
        AtomicLong writerTotalTime = new AtomicLong();
        log.debug("excel writer start");
        try (ExcelWriter excelWriter = EasyExcel.write(outputStream, head).build()) {
            long firstPageSearchStartTime = System.currentTimeMillis();
            Page firstPage = pageFunction.apply(new Page(1, 1000)); // 查首页
            readTotalTime.addAndGet(System.currentTimeMillis() - firstPageSearchStartTime);
            log.debug("[读]：pageNo:{}, total:{}, pageSize:{}, totalPage:{}, pageSearchTime:{}ms", firstPage.getCurrent(), firstPage.getTotal(), firstPage.getSize(), firstPage.getPages(), System.currentTimeMillis() - firstPageSearchStartTime);
            // 如果小于2页，则直接写入，没必要并发
            WriteSheet writeSheet = EasyExcel.writerSheet(0, "Sheet0").build();
            if (firstPage.getPages() <= 2) {
                long firstPageExcelWriteStartTime = System.currentTimeMillis();
                excelWriter.write(firstPage.getRecords(), writeSheet);
                writerTotalTime.addAndGet(System.currentTimeMillis() - firstPageExcelWriteStartTime);
                log.debug("[写]：pageNo:{}, pageExcelWriterTime:{}ms", firstPage.getCurrent(), System.currentTimeMillis() - firstPageExcelWriteStartTime);
                if (firstPage.hasNext()) {
                    long pageSearchStartTime = System.currentTimeMillis();
                    Page page = pageFunction.apply(new Page(2, firstPage.getSize(), false));
                    page.setTotal(firstPage.getTotal());
                    readTotalTime.addAndGet(System.currentTimeMillis() - pageSearchStartTime);
                    log.debug("[读]：pageNo:{}, total:{}, pageSize:{}, totalPage:{}, pageSearchTime:{}ms", page.getCurrent(), page.getTotal(), page.getSize(), page.getPages(), System.currentTimeMillis() - pageSearchStartTime);
                    long pageExcelWriteStartTime = System.currentTimeMillis();
                    excelWriter.write(firstPage.getRecords(), writeSheet);
                    writerTotalTime.addAndGet(System.currentTimeMillis() - pageExcelWriteStartTime);
                    log.debug("[写]：pageNo:{}, pageExcelWriterTime:{}ms", page.getCurrent(), System.currentTimeMillis() - pageExcelWriteStartTime);
                }
                return;
            }
            // 如果大于2页，则并行分页查询，并放到队列中
            int queueCapacity = Math.min(queueSize, (int) firstPage.getPages() - 1); // 如果剩余页数小于设置的队列容量，则队列容量等于剩余页数
            ArrayBlockingQueue<Page> queue = new ArrayBlockingQueue<>(queueCapacity, true); // 用于存放待写入数据
            ThreadPoolExecutor threadPoolExecutor = new ThreadPoolExecutor(1, queueCapacity, 10, TimeUnit.SECONDS, new SynchronousQueue(true), new ThreadPoolExecutor.CallerRunsPolicy());
            // 查询数据使用新线程
            new Thread(() -> {
                try {
                    queue.put(firstPage);
                    // 从第二页开始查到末页，并放到队列中
                    List<CompletableFuture<Page>> pageFutureList = new ArrayList<>(queueCapacity);
                    for (int pageNo = 2; pageNo <= firstPage.getPages(); pageNo++) {
                        int finalPageNo = pageNo;
                        CompletableFuture<Page> pageFuture = CompletableFuture.supplyAsync(() -> {
                            long pageSearchStartTime = System.currentTimeMillis();
                            Page page = pageFunction.apply(new Page(finalPageNo, firstPage.getSize(), false));
                            page.setTotal(firstPage.getTotal());
                            readTotalTime.addAndGet(System.currentTimeMillis() - pageSearchStartTime);
                            log.debug("[读]：pageNo:{}, total:{}, pageSize:{}, totalPage:{}, ActiveCount:{}, pageSearchTime:{}ms", page.getCurrent(), page.getTotal(), page.getSize(), page.getPages(), threadPoolExecutor.getActiveCount(), System.currentTimeMillis() - pageSearchStartTime);
                            return page;
                        }, threadPoolExecutor);
                        pageFutureList.add(pageFuture);
                        // 放入队列
                        if (pageFutureList.size() >= queueCapacity) {
                            for (CompletableFuture<Page> future : pageFutureList) {
                                queue.put(future.join());
                            }
                            pageFutureList.clear();
                        }
                    }
                    // 放入队列
                    for (CompletableFuture<Page> future : pageFutureList) {
                        queue.put(future.join());
                    }
                    pageFutureList.clear();
                } catch (Exception e) {
                    log.error("pageExcelWriterQuick错误！", e);
                    throw new RuntimeException(e);
                }
            }).start();
            // 写入excel，从队列中取数据
            Page page;
            while ((page = queue.poll(PAGE_SEARCH_TIMEOUT_SECONDS, TimeUnit.SECONDS)) != null) {
                long pageExcelWriteStartTime = System.currentTimeMillis();
                int sheetNo = (int) (page.getCurrent() * page.getSize() / EXCEL_SHEET_ROW_MAX_SIZE + 1);
                if (writeSheet.getSheetNo() != sheetNo) {
                    writeSheet = EasyExcel.writerSheet(sheetNo, "Sheet" + sheetNo).build();
                }
                excelWriter.write(page.getRecords(), writeSheet);
                writerTotalTime.addAndGet(System.currentTimeMillis() - pageExcelWriteStartTime);
                log.debug("[写]：pageNo:{}, sheetNo:{}, initQueueCapacity:{}, currentQueueSize:{}, isLastPage:{}, pageExcelWriterTime:{}ms", page.getCurrent(), sheetNo, queueCapacity, queue.size(), !page.hasNext(), System.currentTimeMillis() - pageExcelWriteStartTime);
                if (!page.hasNext()) {break;}
            }
        } catch (Exception e) {
            throw new RuntimeException(e);
        } finally {
            log.debug("excel writer done, totalTime:{}ms, readTotalTime:{}ms, writeTotalTime:{}ms", System.currentTimeMillis() - start, readTotalTime.get(), writerTotalTime.get());
        }
    }
}